15  O Mecanismo de Atenção (Self-Attention)

O mecanismo de Self-Attention (Auto-Atenção) é o componente fundamental que distingue a arquitetura Transformer das redes recorrentes (RNNs) e convolucionais (CNNs) anteriores. Enquanto as RNNs processam sequências passo a passo (o que dificulta a paralelização e a retenção de dependências de longo prazo), o Self-Attention permite que o modelo relacione cada palavra na sequência de entrada com todas as outras palavras simultaneamente.

Este capítulo disseca a Scaled Dot-Product Attention, a fórmula matemática que rege essa operação, e implementa o mascaramento necessário para modelos de linguagem generativos (como a família GPT).

15.1 1. Conceitos Fundamentais: Query, Key e Value

Para calcular a atenção, projetamos a entrada (embeddings) em três vetores distintos para cada token. Esta abordagem é inspirada em sistemas de recuperação de informação:

  1. Query (\(Q\) - Consulta): O que o token atual está procurando? (Representa o foco atual).
  2. Key (\(K\) - Chave): O que o token atual oferece? (Serve como um identificador para correspondência).
  3. Value (\(V\) - Valor): Qual é o conteúdo real da informação? (A informação que será agregada se a chave corresponder à consulta).

Em um mecanismo de Self-Attention, \(Q\), \(K\) e \(V\) originam-se da mesma fonte (a saída da camada anterior), multiplicados por matrizes de pesos aprendíveis (\(W^Q, W^K, W^V\)).

15.2 2. A Fórmula Matemática

A atenção é calculada como uma soma ponderada dos valores (\(V\)), onde o peso atribuído a cada valor é determinado pela compatibilidade entre a consulta (\(Q\)) e a chave correspondente (\(K\)).

\[ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V \]

15.2.1 Passo a Passo do Cálculo:

  1. MatMul (\(QK^T\)): Calcula o produto escalar entre a Query e a Key transposta. Isso resulta em uma Matriz de Scores (ou afinidade). Se o produto escalar é alto, os vetores estão alinhados, indicando alta relevância entre os tokens.
  2. Scale (\(\frac{1}{\sqrt{d_k}}\)): Dividimos os scores pela raiz quadrada da dimensão das chaves (\(d_k\)).
    • Por que escalar? Sem isso, para valores altos de \(d_k\), o produto escalar cresce em magnitude, empurrando a função Softmax para regiões onde os gradientes são extremamente pequenos (vanishing gradients), dificultando o treinamento.
  3. Mask (Opcional): Aplica-se uma máscara (substituindo valores por \(-\infty\)) para impedir que o modelo “veja” certos tokens (ex: tokens futuros ou padding).
  4. Softmax: Converte os scores em probabilidades (pesos de atenção) que somam 1.
  5. **MatMul ($ V\():** Multiplica-se os pesos pelos Valores (\)V$). Isso filtra as informações irrelevantes e preserva as relevantes.

15.3 3. Fluxo de Dados Visual

O diagrama abaixo ilustra o fluxo de tensores dentro do bloco de atenção escalar.

graph TD
    subgraph Inputs
    Q[Query (Q)]
    K[Key (K)]
    V[Value (V)]
    end

    Q --> MatMul1[MatMul (Q x K^T)]
    K --> MatMul1
    
    MatMul1 --> Scale[Scale (1 / sqrt(d_k))]
    
    Scale --> MaskNode{Aplicar Máscara?}
    Mask[Máscara Causal / Padding] -.-> MaskNode
    
    MaskNode -- Sim --> MaskFill[Fill com -inf]
    MaskNode -- Não --> Softmax
    
    MaskFill --> Softmax[Softmax]
    
    Softmax --> MatMul2[MatMul (Weights x V)]
    V --> MatMul2
    
    MatMul2 --> Output[Output Contextualizado]

    style MatMul1 fill:#f9f,stroke:#333,stroke-width:2px
    style Softmax fill:#ff9,stroke:#333,stroke-width:2px
    style MatMul2 fill:#f9f,stroke:#333,stroke-width:2px

15.4 4. Mascaramento (Masking)

Existem dois tipos críticos de máscaras no Transformer:

  1. Padding Mask: Ignora tokens de preenchimento (pad tokens) usados para igualar o tamanho das sentenças em um batch.
  2. Look-ahead (Causal) Mask: Essencial para decodificadores (como no GPT). Ao prever o token na posição \(t\), o modelo não pode ter acesso aos tokens nas posições \(t+1, t+2, \dots\).

A máscara causal é uma matriz triangular superior (excluindo a diagonal) preenchida com \(-\infty\). Quando somada aos scores antes do Softmax: \[ \text{softmax}(x + (-\infty)) = 0 \] Isso garante que a atenção para tokens futuros seja zero.

15.5 5. Implementação em PyTorch

Abaixo, implementamos uma classe robusta para a Scaled Dot-Product Attention, capaz de lidar com máscaras causais.

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class ScaledDotProductAttention(nn.Module):
    """
    Calcula a atenção escalonada por produto escalar.
    Fórmula: Softmax(QK^T / sqrt(d_k)) * V
    """
    
    def __init__(self, dropout_rate=0.1):
        super(ScaledDotProductAttention, self).__init__()
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, query, key, value, mask=None):
        """
        Args:
            query: Tensor de forma (Batch, Heads, Seq_Len_Q, Head_Dim)
            key:   Tensor de forma (Batch, Heads, Seq_Len_K, Head_Dim)
            value: Tensor de forma (Batch, Heads, Seq_Len_K, Head_Dim)
            mask:  Tensor de máscara (Batch, 1, Seq_Len_Q, Seq_Len_K) ou similar.
                   Valores True indicam posições a serem mascaradas (ou 1s/0s dependendo da lógica).
        
        Returns:
            context: Tensor de saída ponderado
            attention_weights: Pesos de atenção (para visualização)
        """
        
        # 1. Obter dimensões
        d_k = query.size(-1)
        
        # 2. Calcular Scores (Q * K^T)
        # Transpomos as duas últimas dimensões de K para permitir a multiplicação
        scores = torch.matmul(query, key.transpose(-2, -1))
        
        # 3. Escalonamento (Scaling)
        scores = scores / math.sqrt(d_k)
        
        # 4. Aplicação da Máscara (Se fornecida)
        if mask is not None:
            # Substitui valores onde a máscara é 0 (ou True, dependendo da implementação) por -inf
            # Aqui assumimos que mask=0 significa "esconder"
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # 5. Softmax para obter probabilidades
        # Aplicado na última dimensão (dimensão da sequência alvo)
        attention_weights = F.softmax(scores, dim=-1)
        
        # 6. Dropout (Opcional, mas comum para regularização)
        attention_weights = self.dropout(attention_weights)
        
        # 7. Multiplicação pelos Valores (Weights * V)
        context = torch.matmul(attention_weights, value)
        
        return context, attention_weights

# --- Exemplo de Uso com Máscara Causal ---

def create_causal_mask(seq_len):
    """Cria uma máscara triangular inferior para impedir 'olhar para o futuro'."""
    # tril retorna a parte triangular inferior da matriz
    mask = torch.tril(torch.ones(seq_len, seq_len))
    return mask

# Configuração de teste
batch_size = 1
heads = 1
seq_len = 4
d_k = 8

# Dados fictícios
q = torch.randn(batch_size, heads, seq_len, d_k)
k = torch.randn(batch_size, heads, seq_len, d_k)
v = torch.randn(batch_size, heads, seq_len, d_k)

# Criar máscara causal (1s visíveis, 0s ocultos)
causal_mask = create_causal_mask(seq_len) 
# Expande dimensões para broadcasting: (1, 1, seq_len, seq_len)
causal_mask = causal_mask.unsqueeze(0).unsqueeze(0)

# Instanciar e rodar
attention_layer = ScaledDotProductAttention()
output, weights = attention_layer(q, k, v, mask=causal_mask)

print("Pesos de Atenção (com máscara causal):")
print(weights[0][0]) 
# Note que a parte superior direita será 0.

15.5.1 Análise da Saída do Código

Se executarmos o código acima, a matriz de pesos de atenção (weights) terá uma estrutura triangular inferior. Por exemplo, para uma sequência de tamanho 4:

  • O token 1 só atende ao token 1.
  • O token 2 atende aos tokens 1 e 2.
  • O token 3 atende aos tokens 1, 2 e 3.
  • As posições onde \(j > i\) (futuro) terão peso \(0\) (após o softmax de \(-\infty\)).

15.6 Conclusão do Capítulo

O mecanismo de Scaled Dot-Product Attention é eficiente computacionalmente, pois se baseia em operações de matriz altamente otimizadas. A introdução do fator de escala \(\sqrt{d_k}\) garante estabilidade numérica, e o uso de máscaras permite que a mesma arquitetura seja utilizada tanto para codificação (entender todo o contexto) quanto para decodificação (gerar texto sequencialmente). No próximo capítulo, veremos como expandir este conceito para Multi-Head Attention, permitindo que o modelo foque em diferentes subespaços de representação simultaneamente.